import imaplib
import smtplib
import socket
import socks
import ssl
import time
from api.utils.logger import get_logger
from api.utils.singbox_proxy_helper import get_singbox_proxy_config

logger = get_logger("proxy_helper")
sync_logger = get_logger("sync")

# --- 核心修复：自定义连接类 ---

class ProxyIMAP4_SSL(imaplib.IMAP4_SSL):
    """支持指定 SOCKS5 代理的 IMAP_SSL 类"""
    def __init__(self, host, port, proxy_setting, keyfile=None, certfile=None, ssl_context=None, timeout=None):
        self.proxy_setting = proxy_setting
        if ssl_context and keyfile:
            raise ValueError("ssl_context and keyfile arguments are mutually exclusive")
        if ssl_context and certfile:
            raise ValueError("ssl_context and certfile arguments are mutually exclusive")

        self.keyfile = keyfile
        self.certfile = certfile
        self.ssl_context = ssl_context
        self.sslobj = None
        imaplib.IMAP4.__init__(self, host, port, timeout)

    def _create_socket(self, timeout=None):
        sock = socks.socksocket()
        sock.set_proxy(
            socks.SOCKS5,
            self.proxy_setting['host'],
            self.proxy_setting['port'],
            username=self.proxy_setting.get('username'),
            password=self.proxy_setting.get('password')
        )
        
        effective_timeout = timeout if timeout is not None else self.timeout
        sock.settimeout(effective_timeout)
        
        try:
            sock.connect((self.host, self.port))
        except socks.SOCKS5Error as e:
            sync_logger.error(f"[IMAP] SOCKS5 代理认证或协商失败: {e}")
            raise
        except socks.ProxyConnectionError as e:
            sync_logger.error(f"[IMAP] 代理连接失败: 无法连接到本地代理服务 ({self.proxy_setting['host']}:{self.proxy_setting['port']})。请确认代理程序（如 Sing-box）正在运行。错误: {e}")
            raise
        except socks.GeneralProxyError as e:
            sync_logger.error(f"[IMAP] 通用代理错误 (通常意味着代理无法连接到目标主机): {e}")
            raise
            
        if self.ssl_context:
            ssock = self.ssl_context.wrap_socket(sock, server_hostname=self.host)
        else:
            ctx = ssl.create_default_context()
            ssock = ctx.wrap_socket(sock, server_hostname=self.host)
        return ssock

class ProxySMTP_SSL(smtplib.SMTP_SSL):
    """支持 SOCKS5 代理的 SMTP_SSL 类 (用于 465 端口)"""
    def __init__(self, host, port, proxy_config, **kwargs):
        self.proxy_config = proxy_config
        super().__init__(host='', port=port, **kwargs)
        self.host = host
        self.port = port
        
        try:
            self._custom_connect()
        except Exception:
            self.close()
            raise

    def _custom_connect(self):
        sync_logger.debug(f"[SMTP_SSL] 正在通过代理 {self.proxy_config['host']}:{self.proxy_config['port']} 连接 {self.host}:{self.port}...")
        
        self.sock = socks.socksocket()
        self.sock.set_proxy(
            socks.SOCKS5,
            self.proxy_config['host'],
            self.proxy_config['port'],
            username=self.proxy_config.get('username'),
            password=self.proxy_config.get('password')
        )
        self.sock.settimeout(self.timeout or 30)

        try:
            self.sock.connect((self.host, self.port))
            sync_logger.debug("[SMTP_SSL] SOCKS5 握手成功，TCP 通道已建立。")
        except socks.ProxyConnectionError as e:
            sync_logger.error(f"[SMTP_SSL] 代理连接失败: Sing-box 拒绝连接 - {e}")
            raise
        except socks.GeneralProxyError as e:
            sync_logger.error(f"[SMTP_SSL] 代理错误: Sing-box 无法连接远程服务器 - {e}")
            raise

        sync_logger.debug("[SMTP_SSL] 正在进行 SSL 握手...")
        try:
            self.sock = self.context.wrap_socket(
                self.sock, 
                server_hostname=self.host 
            )
            self.file = self.sock.makefile('rb')
            sync_logger.debug("[SMTP_SSL] SSL 握手成功！")
        except ssl.SSLEOFError:
            sync_logger.error("[SMTP_SSL] SSL EOF 错误: 远程服务器关闭了连接。可能原因：SNI缺失、IP被黑名单、端口错误。")
            raise
        except ssl.SSLError as e:
            sync_logger.error(f"[SMTP_SSL] SSL 协议错误: {e}")
            raise

        (code, msg) = self.getreply()
        sync_logger.debug(f"[SMTP_SSL] 服务器响应: {code} {msg}")
        if code != 220:
            self.close()
            raise smtplib.SMTPConnectError(code, msg)

class ProxySMTP_STARTTLS(smtplib.SMTP):
    """支持 SOCKS5 的 SMTP 类 (用于 25/587 端口的 STARTTLS)"""
    def __init__(self, host, port, proxy_config, **kwargs):
        self.proxy_config = proxy_config
        super().__init__(host, port, **kwargs)

    def _get_socket(self, host, port, timeout):
        sync_logger.debug(f"[SMTP_STARTTLS] 正在通过代理连接 {host}:{port}...")
        sock = socks.socksocket()
        sock.set_proxy(socks.SOCKS5, self.proxy_config['host'], self.proxy_config['port'])
        sock.settimeout(timeout)
        sock.connect((host, port))
        return sock

# --- 辅助逻辑 ---
def _prepare_proxy_info(proxy_config: dict) -> dict:
    proxy_type = proxy_config.get('type', 'socks5').lower()
    if proxy_type in ['vmess', 'vless', 'trojan', 'shadowsocks', 'hysteria', 'hysteria2', 'tuic', 'wireguard']:
        v2ray_config = proxy_config.get('v2ray_config')
        bridge = get_singbox_proxy_config(v2ray_config)
        if bridge.get('type') == 'unsupported':
             raise ConnectionError(f"代理启动失败: {bridge.get('error')}")
        sync_logger.info(f"Sing-box 桥接成功，本地SOCKS地址: {bridge['host']}:{bridge['port']}")
        return {'host': bridge['host'], 'port': bridge['port'], 'username': None, 'password': None}
    else:
        sync_logger.debug(f"使用标准代理: {proxy_type}")
        return {'host': proxy_config['host'], 'port': proxy_config['port'], 'username': proxy_config.get('username'), 'password': proxy_config.get('password')}

def _create_loose_ssl_context():
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx

# --- 主调用函数 ---
def create_imap_connection_with_proxy(server: str, port: int, use_ssl: bool, proxy_config: dict, timeout: int = 60):
    try:
        proxy_setting = _prepare_proxy_info(proxy_config)
        sync_logger.debug(f"[IMAP] 准备连接到 {server}:{port}，使用代理: {proxy_setting['host']}:{proxy_setting['port']}")
        if use_ssl:
            conn = ProxyIMAP4_SSL(server, port, proxy_setting=proxy_setting, ssl_context=_create_loose_ssl_context(), timeout=timeout)
        else:
            sync_logger.warning(f"检测到非 SSL IMAP 连接 (服务器: {server}:{port})。此方式不推荐。")
            _original_socket = socket.socket
            try:
                socks.set_default_proxy(socks.SOCKS5, proxy_setting['host'], proxy_setting['port'])
                socket.socket = socks.socksocket
                conn = imaplib.IMAP4(server, port, timeout=timeout)
            finally:
                socket.socket = _original_socket
                socks.set_default_proxy()
        sync_logger.info(f"[IMAP] 连接成功: {server}:{port}")
        return conn
    except imaplib.IMAP4.abort:
        raise
    except (socks.GeneralProxyError, socks.ProxyConnectionError, socks.SOCKS5Error, TimeoutError, ConnectionRefusedError, OSError) as e:
        raise
    except Exception as e:
        sync_logger.error(f"IMAP 连接时发生未知错误: {e}", exc_info=True)
        raise

def create_smtp_connection_with_proxy(server: str, port: int, use_ssl: bool, proxy_config: dict, timeout: int = 60):
    try:
        proxy_setting = _prepare_proxy_info(proxy_config)
        sync_logger.debug(f"[SMTP] 准备连接到 {server}:{port}，使用代理: {proxy_setting['host']}:{proxy_setting['port']}")
        
        if use_ssl and port == 465:
            conn = ProxySMTP_SSL(server, port, proxy_config=proxy_setting, context=_create_loose_ssl_context(), timeout=timeout)
        elif not use_ssl and port == 587:
            conn = ProxySMTP_STARTTLS(server, port, proxy_config=proxy_setting, timeout=timeout)
            sync_logger.debug(f"[SMTP] 端口为 587，尝试 STARTTLS...")
            try:
                conn.starttls(context=_create_loose_ssl_context())
                sync_logger.debug(f"[SMTP] STARTTLS 成功。")
            except Exception as e:
                sync_logger.error(f"SMTP STARTTLS 失败: {e}", exc_info=True)
                conn.close()
                raise
        else:
            # 兼容其他 use_ssl=False 且非 587 的情况，或 use_ssl=True 但非 465 的情况
            # 这种组合比较模糊，我们优先使用普通 SMTP 类，如果需要 STARTTLS，调用者需自行处理
            sync_logger.warning(f"SMTP 连接端口({port})与SSL设置({use_ssl})组合不标准，使用通用SMTP连接。")
            conn = ProxySMTP_STARTTLS(server, port, proxy_config=proxy_setting, timeout=timeout)

        sync_logger.info(f"[SMTP] 连接成功: {server}:{port}")
        return conn
    except (smtplib.SMTPException, socks.GeneralProxyError, socks.ProxyConnectionError, socks.SOCKS5Error, TimeoutError, ConnectionRefusedError, OSError) as e:
        raise
    except Exception as e:
        raise